package com.example;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class SQLConditionAdder {
    
    /**
     * 添加条件到SELECT语句中。该方法会解析传入的SQL语句，并在指定的表名上添加条件。
     * 如果SQL语句是SELECT语句，则会遍历其FROM子句和子查询，找到匹配的表名并添加条件。
     * 如果表名匹配且WHERE子句已存在，则使用AND操作符将新条件与现有条件连接；
     * 如果WHERE子句不存在，则直接设置新条件。
     *
     * @param sql 需要修改的SQL语句，必须是合法的SQL语句
     * @param tableName 需要添加条件的表名
     * @param condition 需要添加的条件，格式为"字段名=值"，例如"creator=0001"
     * @return 修改后的SQL语句，如果SQL语句不是SELECT语句，则返回原SQL语句
     * @throws JSQLParserException 如果SQL语句解析失败，抛出此异常
     */
    public static String addConditionToSelect(String sql, String tableName, String condition) throws JSQLParserException {
        Statement statement = CCJSqlParserUtil.parse(sql);
        if (statement instanceof Select) {
            Select select = (Select) statement;
            SelectBody selectBody = select.getSelectBody();
            processSelectBody(selectBody,tableName, condition);
            return select.toString();
        }
        return sql;
    }

    private static void processSelectBody(SelectBody selectBody, String tableName, String condition) {
        if (selectBody instanceof PlainSelect) {
            PlainSelect plainSelect = (PlainSelect) selectBody;
            
	        if (plainSelect.getFromItem() instanceof Table) {
                Table table = (Table) plainSelect.getFromItem();
                if (tableName.contains(table.getName())) {
                    String tableAlias = table.getAlias() == null ? "" : table.getAlias().getName();
                    if (plainSelect.getWhere() != null) {
                        plainSelect.setWhere(new AndExpression(plainSelect.getWhere(), createCondition(condition, tableAlias)));
                    } else {
                        plainSelect.setWhere(createCondition(condition, tableAlias));
                    }
                }
	        }
            // 处理子查询
            if (plainSelect.getFromItem() instanceof SubSelect) {
                SubSelect subSelect = (SubSelect) plainSelect.getFromItem();
                processSelectBody(subSelect.getSelectBody(), tableName,condition);
            }
            for (SelectItem selectItem : plainSelect.getSelectItems()) {
                if (selectItem instanceof SelectExpressionItem) {
                    SelectExpressionItem selectExpressionItem = (SelectExpressionItem) selectItem;
                    if (selectExpressionItem.getExpression() instanceof SubSelect) {
                        SubSelect subSelect = (SubSelect) selectExpressionItem.getExpression();
                        processSelectBody(subSelect.getSelectBody(), tableName, condition);
                    }
                }
            }
            // 处理子查询类型的条件值
            if (plainSelect.getWhere() != null) {
                plainSelect.getWhere().accept(new ExpressionVisitorAdapter() {
                    @Override
                    public void visit(InExpression inExpression) {
                        if (inExpression.getRightItemsList() instanceof SubSelect) {
                            SubSelect subSelect = (SubSelect) inExpression.getRightItemsList();
                            processSelectBody(subSelect.getSelectBody(), tableName, condition);
                        }
                        super.visit(inExpression);
                    }

                    @Override
                    public void visit(EqualsTo equalsTo) {
                        if (equalsTo.getRightExpression() instanceof SubSelect) {
                            SubSelect subSelect = (SubSelect) equalsTo.getRightExpression();
                            processSelectBody(subSelect.getSelectBody(), tableName, condition);
                        }
                        super.visit(equalsTo);
                    }
                });
            }
        }
    }

    private static Expression createCondition(String condition, String tableAlias) {
	    try {
            Expression expression = CCJSqlParserUtil.parseCondExpression(condition);
            expression.accept(new ExpressionVisitorAdapter() {
                @Override
                public void visit(Column column) {
                    if (column.getTable() == null) {
                        column.setTable(new Table(tableAlias));
                    }
                    super.visit(column);
                }
            });
            return expression;
        } catch (JSQLParserException e) {
            return new EqualsTo();
	    }
    }
    
    public static void main(String[] args) {
        List<String> testCases = Arrays.asList(
                // 简单 SELECT 语句
                "SELECT * FROM orders",
                
                // 带 WHERE 条件的 SELECT 语句
                "SELECT * FROM orders WHERE order_date > '2024-01-01'",
                
                // 多表 JOIN 语句
                "SELECT o.order_id, c.customer_name FROM orders o JOIN customers c ON o.customer_id = c.customer_id",
                
                // 含有子查询的语句
                "SELECT customer_id, (SELECT SUM(amount) FROM orders WHERE orders.customer_id = customers.customer_id) AS total_amount FROM customers",
                
                // GROUP BY 语句
                "SELECT customer_id, COUNT(*) FROM orders GROUP BY customer_id",
                
                // HAVING 语句
                "SELECT customer_id, COUNT(*) as order_count FROM orders GROUP BY customer_id HAVING order_count > 5",
                
                // 嵌套子查询
                "SELECT * FROM customers WHERE customer_id IN (SELECT customer_id FROM orders WHERE amount > 1000)",
                
                // 复杂查询
                "SELECT c.customer_id, c.customer_name, (SELECT SUM(o.amount) FROM orders o WHERE o.customer_id = c.customer_id) AS total_spent FROM customers c"
        );
        
        for (String sql : testCases) {
            // String modifiedSql = SQLConditionAdder.addConditionToSelect(sql, tableName, condition);
            String modifiedSql = SQLConditionAdder.addConditionToSelect(sql, new HashMap<String, String>() {{
                put("orders", "creator is null or creator in (0001,0003)");
                put("customers", "creator is null or creator in (0002,0004)");
            }});
            System.out.println("\n原始SQL:");
            System.out.println(sql);
            System.out.println("\n修改后SQL:");
            System.out.println(modifiedSql);
            System.out.println("--------------------------------------------------------------");
        }
        
    }
    
    static String addConditionToSelect(String sql, Map<String, String> map){
        String newSql = sql;
        for (Map.Entry<String, String> entry : map.entrySet()) {
            String tableName = entry.getKey();
            String condition = entry.getValue();
            try {
                newSql = SQLConditionAdder.addConditionToSelect(newSql, tableName, condition);
            } catch (JSQLParserException e) {
                e.printStackTrace();
            }
        }
        return newSql;
    }
}
